import streamlit as st
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from torch.utils.data import TensorDataset, DataLoader

# ========== 1. 数据加载与预处理 ==========
@st.cache_data
def load_and_preprocess_data():
    # 读取CSV文件
    df = pd.read_csv("iris.csv")

    # 特征和标签
    X = df.drop("species", axis=1).values
    y_str = df["species"].values

    # 手动映射标签
    label_map = {
        "Iris-setosa": 0,
        "Iris-versicolor": 1,
        "Iris-virginica": 2
    }
    y = [label_map[label] for label in y_str]

    # 标准化特征
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    # 转换为tensor
    X_tensor = torch.tensor(X, dtype=torch.float32)
    y_tensor = torch.tensor(y, dtype=torch.long)

    return X_tensor, y_tensor, df.columns[:-1], label_map

# ========== 2. 定义模型 ==========
class IrisNet(nn.Module):
    def __init__(self):
        super(IrisNet, self).__init__()
        self.fc1 = nn.Linear(4, 16)
        self.fc2 = nn.Linear(16, 3)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# ========== 3. 训练函数 ==========
def train_model(X, y, epochs=50, batch_size=16, lr=0.01):
    dataset = TensorDataset(X, y)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = IrisNet()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)

    loss_list = []
    for epoch in range(epochs):
        total_loss = 0
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        loss_list.append(total_loss / len(dataloader))

    # 保存模型
    torch.save(model.state_dict(), "iris_model.pth")
    return model, loss_list

# ========== 4. Streamlit 主界面 ==========
st.title("🌸 鸢尾花分类（Iris Classification）")
st.write("手动映射标签，不依赖 sklearn 的 LabelEncoder")

# 加载数据
X, y, feature_names, label_map = load_and_preprocess_data()
target_names = {v: k for k, v in label_map.items()}  # 反向映射

# 分割训练测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 按钮触发训练
if st.button("开始训练模型"):
    model, loss_list = train_model(X_train, y_train, epochs=100)
    st.success("✅ 模型训练完成并保存为 iris_model.pth")

    # 损失曲线
    st.line_chart(loss_list)

# ========== 5. 推理部分 ==========
st.subheader("🔮 模型推理")
st.write("输入鸢尾花的四个特征，模型将预测类别")

col1, col2 = st.columns(2)
with col1:
    sepal_length = st.number_input("Sepal Length (花萼长度)", min_value=0.0, max_value=10.0, value=5.1)
    sepal_width = st.number_input("Sepal Width (花萼宽度)", min_value=0.0, max_value=10.0, value=3.5)
with col2:
    petal_length = st.number_input("Petal Length (花瓣长度)", min_value=0.0, max_value=10.0, value=1.4)
    petal_width = st.number_input("Petal Width (花瓣宽度)", min_value=0.0, max_value=10.0, value=0.2)

if st.button("预测类别"):
    # 加载模型
    model = IrisNet()
    model.load_state_dict(torch.load("iris_model.pth"))
    model.eval()

    input_tensor = torch.tensor([[sepal_length, sepal_width, petal_length, petal_width]], dtype=torch.float32)

    # 推理
    with torch.no_grad():
        output = model(input_tensor)
        predicted_class = torch.argmax(output, dim=1).item()
    st.write(predicted_class)
    st.write(f"{target_names[predicted_class]}")
    st.write(f"🌼 预测结果： **{target_names[predicted_class]}**")
